import os 
import signal
import sys 
import time 
import threading
import asyncio 
import websockets 
import json
import numpy as np
from concurrent.futures  import ThreadPoolExecutor
from threading import Lock, Event 
from funasr import AutoModel
import logging
# print(websockets.__version__) 15.0.1
# 使用魔塔的paraformer-zh-streaming模型
# 配置日志 
logging.basicConfig(level=logging.INFO) 
logger = logging.getLogger(__name__) 
 
# 全局配置 
CONFIG = {
    'max_connections': 100,
    'worker_threads': 8,
    'heartbeat_timeout': 20,
    'sample_rate': 16000,
    'format_pcm': 'pcm',
    # 魔塔模型参数
    'chunk_size': [0, 14, 5], # [0, 10, 5] 600ms, [0, 8, 4] 480ms，[0, 14, 5]840ms 13440采样点
    'encoder_chunk_look_back': 4, # number of chunks to lookback for encoder self-attention
    'decoder_chunk_look_back': 1  # number of encoder chunks to lookback for decoder cross-attention
}
 
class ClientSession:
    """WebSocket客户端会话管理"""
    def __init__(self, websocket, client_id):
        self.websocket  = websocket 
        self.client_id  = client_id 
        self.last_active  = time.time() 
        self.text_buffer  = ""
        self.lock  = Lock()
        self.active  = True 
        self.recognition  = None 
        self.callback  = None
        self.init_recognition() 

 
    def init_recognition(self):
        """初始化语音识别会话"""
        # 初始化魔塔模型
        self.model = AutoModel(model="paraformer-zh-streaming", model_revision="v2.0.4")
        self.cache = {}
        self.frames_count = 0
        # 计算每个音频块的大小（以样本数为单位）
        self.chunk_stride = CONFIG['chunk_size'][1] * 960  # 600ms
        
        # 不再需要callback和start，魔塔模型直接处理音频数据
        self.recognition = self  # 为了兼容原有代码，将self赋值给recognition
    
    def send_audio_frame(self, audio_data):
        """处理音频数据"""
        try:
            # 将二进制音频数据转换为numpy数组
            audio_array = np.frombuffer(audio_data, dtype=np.int16).astype(np.float32) / 32768.0
            
            # 每隔chunk_stride个样本处理一次
            self.frames_count += len(audio_array)
            if self.frames_count >= self.chunk_stride or len(audio_array) >= self.chunk_stride:
                # 是否为最后一帧（这里设为False，因为是实时流）
                is_final = False
                
                # 调用模型处理音频数据
                res = self.model.generate(
                    input=audio_array, 
                    cache=self.cache, 
                    is_final=is_final, 
                    chunk_size=CONFIG['chunk_size'], 
                    encoder_chunk_look_back=CONFIG['encoder_chunk_look_back'], 
                    decoder_chunk_look_back=CONFIG['decoder_chunk_look_back']
                )
                
                # 处理识别结果
                if res and len(res) > 0 and 'text' in res[0]:
                    with self.lock:
                        text = res[0]['text']
                        if text.strip():
                            self.text_buffer += text
                            # print(f"[Client {self.client_id}] Sentence: {text}")
                
                # 重置计数器
                self.frames_count = 0
        except Exception as e:
            logger.error(f"Error processing audio frame: {str(e)}")
    
    def start(self):
        """兼容原有代码的start方法"""
        pass
    
    def stop(self):
        """兼容原有代码的stop方法"""
        self.cache = {}
 
    def update_activity(self):
        """更新最后活动时间"""
        with self.lock: 
            self.last_active  = time.time() 
 
    async def close(self):
        """安全关闭会话"""
        with self.lock: 
            self.active  = False 
    
            if self.recognition: 
                self.stop() 
            if self.websocket: 
                try:
                    await self.websocket.close() 
                except:
                    pass 
            logger.info(f"[Session]  Closed connection for client {self.client_id}") 
 
class ClientManager:
    """WebSocket客户端连接管理器"""
    def __init__(self):
        self.sessions  = {}
        self.lock  = Lock()
        self.heartbeat_event  = Event()
        self.executor  = ThreadPoolExecutor(max_workers=CONFIG['worker_threads'])
 
    def add_session(self, websocket,client_id):
        """添加新会话"""
        with self.lock: 
            if len(self.sessions)  >= CONFIG['max_connections']:
                logger.warning("Connection  limit reached")
                return None
            
            if client_id in self.sessions:
                print('同一个客户端启用一个会话') 
                session = self.sessions[client_id]
                # session.recognition.start()
                session.websocket  = websocket  # 更新WebSocket连接
                session.active  = True
            else:
                print('新的客户端，建立新会话')
                session = ClientSession(websocket, client_id)
                self.sessions[client_id]  = session
                # logger.info(f"[Manager]  New session {client_id}, total: {len(self.sessions)}") 
            return session

 
    async def remove_session(self, client_id):
        """移除会话"""
        with self.lock: 
            if client_id in self.sessions: 
                # await self.sessions[client_id].close() 
                del self.sessions[client_id] 
                logger.info(f"[Manager]  Removed {client_id}, remaining: {len(self.sessions)}") 
 
    def heartbeat_check(self):
        """心跳检测线程"""
        while not self.heartbeat_event.is_set(): 
            time.sleep(5) 
            now = time.time() 
            print(f"Heartbeat  check at {time.strftime('%Y-%m-%d  %H:%M:%S', time.localtime(now))}")
            logger.info(f"Heartbeat  check at {time.strftime('%Y-%m-%d  %H:%M:%S', time.localtime(now))}") 
            # with self.lock: 
            #     to_remove = []
            #     for client_id, session in self.sessions.items(): 
            #         if now - session.last_active  > CONFIG['heartbeat_timeout']:
            #             logger.warning(f"[Heartbeat]  Timeout {client_id}")
            #             to_remove.append(client_id) 
 
            #     for client_id in to_remove:
            #         asyncio.create_task(self.remove_session(client_id)) 
 
    async def shutdown(self):
        """安全关闭所有连接"""
        self.heartbeat_event.set() 
        with self.lock: 
            for session in self.sessions.values(): 
                await session.close() 
            self.sessions.clear() 
 

 
async def handle_client(websocket, path, manager):
    from urllib.parse  import parse_qs, urlparse
    query = parse_qs(urlparse(path).query)
    client_id = query.get('client_id',  [''])[0]
    print(f"建立连接了{client_id}")
    logger.info(f"链接建立了{client_id}")

    """处理WebSocket客户端连接"""
    session = manager.add_session(websocket,client_id) 
    if not session:
        await websocket.close() 
        return
 
    try:
        async for message in websocket:
            session.update_activity() 
            if isinstance(message, str):
                if message == "!DISCONNECT":
                    print("接收到终止信号-------------")      
                    print(session.text_buffer)
                    logger.info(session.text_buffer)
                    await websocket.send(session.text_buffer) 
                    session.text_buffer  = ''
                    try:
                        await websocket.close() 
                        await manager.remove_session(session.client_id)
                        print('关闭了session和WebSocket')
                    except Exception as e:
                        print(e)
            elif isinstance(message, bytes):
                #print(f"接收到音频数据: {len(message)} 字节")
                # 处理二进制音频数据
                session.send_audio_frame(message)
                
            else:
                logger.warning(f"Received  unexpected message type from {session.client_id}") 
 
    except websockets.exceptions.ConnectionClosed: 
        logger.info(f"[Client  {session.client_id}]  Connection closed")
    except Exception as e:
        logger.error(f"[Client  {session.client_id}]  Error: {str(e)}")
    # finally:
    #     await manager.remove_session(session.client_id) 
 
async def start_server(manager):
    """优化后的跨平台WebSocket服务启动函数"""
    loop = asyncio.get_running_loop() 
    stop = loop.create_future() 
    
    # 信号处理（兼容Windows/Unix）
    if sys.platform  == 'win32':
        signal.signal(signal.SIGINT,  lambda *_: loop.call_soon_threadsafe(stop.set_result,  None))
    else:
        for sig in (signal.SIGINT, signal.SIGTERM):
            loop.add_signal_handler(sig,  stop.set_result,  None)
 
    server_config = {
        "host": "0.0.0.0",
        "port": 12345,
        "reuse_address": True,  # 关键修改点 
        "ping_interval": 30,    # 新增心跳检测 
        "ping_timeout": 60,
        "max_queue": 1024       # 控制连接队列
    }
 
    # 非Windows平台追加配置 
    if sys.platform  != 'win32':
        server_config.update({ 
            "reuse_port": True,  # Unix专用优化
            "backlog": 100       # 连接等待队列
        })
    
    async with websockets.serve( 
        lambda websocket,path: handle_client(websocket,path, manager),
        **server_config
    ):
        print(f"[2025-07-11  14:50] 服务已启动 {server_config['host']}:{server_config['port']}")
        logger.info(f"[2025-07-11  14:50] 服务已启动 {server_config['host']}:{server_config['port']}")
        await stop
        print(" 服务优雅终止中...")
        logger.info(" 服务优雅终止中...")
        # await manager.shutdown() 
 
def signal_handler(sig, frame):
    """处理中断信号"""
    print('\nServer shutting down gracefully...')
    sys.exit(0) 
 
if __name__ == '__main__':
    logger.info("Starting  real-time WebSocket server")
 
    # 初始化管理器 
    manager = ClientManager()
   
    signal.signal(signal.SIGINT,  signal_handler)
   
    # 启动心跳线程 
    threading.Thread(target=manager.heartbeat_check,  daemon=True).start()

    # 启动WebSocket服务器 
    asyncio.run(start_server(manager))
    


